package com.wanglei.accumulator

import org.apache.spark.util.AccumulatorV2
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

// 累加器：分布式共享只写变量
object Spark02_MYacc {

  def main(args: Array[String]): Unit = {

    //
    val conf = new SparkConf().setMaster("local[2]").setAppName("acc")
    val sc = new SparkContext(conf)


    //
    val rdd = sc.makeRDD(List(
      "hello spark", "hello scala"
    )).flatMap(
      line => line.split(" ")
    )

    // 自定义累加器完成wordcount
    val wc = new MyAcc
    sc.register(wc, "myacc")

    rdd.foreach(
      word => wc.add(word)
    )

    println(wc.value)



    //
    sc.stop()

  }

  class MyAcc extends AccumulatorV2[String, mutable.Map[String, Long]] {

    private var wcMap = mutable.Map[String, Long]()

    // 判断累加器是不是初始状态，当map为空时就是初始状态
    override def isZero: Boolean = wcMap.isEmpty

    override def copy(): AccumulatorV2[String, mutable.Map[String, Long]] = new MyAcc

    override def reset(): Unit = wcMap.clear()

    override def add(v: String): Unit = {
      val cnt = wcMap.getOrElse(v, 0L) + 1
      wcMap.update(v, cnt)
    }


    // 在driver合并累加器
    override def merge(other: AccumulatorV2[String, mutable.Map[String, Long]]): Unit = {
      val map1 = this.wcMap
      val map2 = other.value

      map2.foreach {
        case (word, cnt) => {
          val newCnt = map1.getOrElse(word, 0L) + cnt
          map1.update(word, newCnt)
        }
      }
    }

    override def value: mutable.Map[String, Long] = wcMap
  }
}
